{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "from astronet.astro_cnn_model import input_ds\n",
    "from astronet.util import config_util\n",
    "\n",
    "data_files = '/mnt/tess/astronet/tfrecords-new\\+old/test-0000[6-7]*'\n",
    "\n",
    "parent = '/mnt/tess/astronet/checkpoints/local_global_new_old_2/10'\n",
    "all_dirs = os.listdir(parent)\n",
    "d, = all_dirs\n",
    "model_dir = os.path.join(parent, d)\n",
    "\n",
    "model = tf.saved_model.load(model_dir)\n",
    "config = config_util.load_config(model_dir)\n",
    "\n",
    "ds = input_ds.build_dataset(\n",
    "    file_pattern=data_files,\n",
    "    input_config=config.inputs,\n",
    "    batch_size=1,\n",
    "    include_labels=False,\n",
    "    shuffle_filenames=False,\n",
    "    repeat=1,\n",
    "    include_identifiers=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(input_config, feature, val):\n",
    "    mean = config.inputs.features[feature]['mean']\n",
    "    std = config.inputs.features[feature]['std']\n",
    "    return (val - mean) / std\n",
    "\n",
    "\n",
    "def denormalize(input_config, feature, val):\n",
    "    mean = config.inputs.features[feature]['mean']\n",
    "    std = config.inputs.features[feature]['std']\n",
    "    return val * std + mean\n",
    "\n",
    "\n",
    "def sweep_inputs(tic_id, ds, config, feature, min_val, max_val, n, n2):\n",
    "    def override(i, data):\n",
    "        i = tf.cast(i // (n2 + 1), tf.float32)\n",
    "        data[0][feature] = normalize(\n",
    "            config.inputs, feature, \n",
    "            data[0][feature] * 0 + (\n",
    "                min_val + i * (max_val - min_val) / n))\n",
    "        return data\n",
    "    \n",
    "    def select_tic(_, identifiers):\n",
    "        return tf.squeeze(identifiers == tic_id)\n",
    "\n",
    "    ds = ds.filter(select_tic)\n",
    "    ds = ds.repeat(n + 1)\n",
    "    ds = ds.enumerate()\n",
    "    ds = ds.map(override)\n",
    "    \n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import ticker\n",
    "\n",
    "def plot_sweep(sweep_predictions, label, n, features, fig, subplot):\n",
    "    assert len(features) == 2\n",
    "    pred_e = sweep_predictions[label].values.reshape(n + 1, n + 1)\n",
    "    x = sweep_predictions[features[0]].values.reshape(n + 1, n + 1)[0, :]\n",
    "    y = sweep_predictions[features[1]].values.reshape(n + 1, n + 1)[:, 0]\n",
    "\n",
    "    ax = fig.add_subplot(subplot)\n",
    "    im = ax.imshow(pred_e, vmin=0, vmax=1, cmap=plt.get_cmap('RdYlGn'))\n",
    "    ax.set_xticklabels([''] + ['%3.3f' % v for v in x])\n",
    "    ax.set_yticklabels([''] + ['%3.3f' % v for v in y])\n",
    "    ax.set_xlabel(features[0])\n",
    "    ax.set_ylabel(features[1])\n",
    "    plt.title(label)\n",
    "    return im\n",
    "    \n",
    "def sweep(tic_id, ds, config, n, featurs_and_ranges):\n",
    "    fr_1, fr_2 = featurs_and_ranges\n",
    "    \n",
    "    name_1, mn, mx = fr_1\n",
    "    ds = sweep_inputs(tic_id, ds, config, name_1, mn, mx, n, 0)\n",
    "    name_2, mn, mx = fr_2\n",
    "    ds = sweep_inputs(tic_id, ds, config, name_2, mn, mx, n, n)\n",
    "    \n",
    "    labels = [\"disp_E\", \"disp_N\", \"disp_J\", \"disp_S\", \"disp_B\"]\n",
    "\n",
    "    data = []\n",
    "    for features, identifiers in ds:\n",
    "        preds = model(features)\n",
    "        row = dict(zip(labels, preds.numpy()[0]))\n",
    "        row[name_1] = denormalize(\n",
    "            config, name_1, features[name_1].numpy().item())\n",
    "        row[name_2] = denormalize(\n",
    "            config, name_2, features[name_2].numpy().item())\n",
    "        data.append(row)\n",
    "\n",
    "    sweep_predictions = pd.DataFrame(data)\n",
    "    \n",
    "    fig = plt.figure(figsize=(20, 14))\n",
    "    subplot = 230\n",
    "    for l in labels:\n",
    "        subplot += 1\n",
    "        im = plot_sweep(sweep_predictions, l, n, (name_1, name_2), fig, subplot)\n",
    "    fig.add_subplot(subplot + 1)\n",
    "    fig.colorbar(im)\n",
    "    \n",
    "    return sweep_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sweep_predictions = sweep(\n",
    "    293665214, ds, config, 4,\n",
    "    (('star_mass', 0.14, 4), ('star_rad', 0.17, 12)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sweep_predictions = sweep(\n",
    "    160991022, ds, config, 4,\n",
    "    (('star_rad', 0.17, 50), ('Period', 0.02, 50)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_ds = ds.filter(lambda _, ids: tf.squeeze(ids == 293665214))\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy import ndimage\n",
    "\n",
    "for features, _ in filtered_ds:\n",
    "    fig = plt.figure(figsize=(12, 3))\n",
    "    fig.add_subplot(131)\n",
    "    gv = features['global_view'][0].numpy()\n",
    "    plt.plot(gv)\n",
    "    gv = ndimage.gaussian_filter(gv, 0.1)\n",
    "    plt.plot(gv)\n",
    "    \n",
    "    fig.add_subplot(132)\n",
    "    lv = features['local_view'][0].numpy()\n",
    "    plt.plot(lv)\n",
    "    lv = ndimage.gaussian_filter1d(lv, 2)\n",
    "    plt.plot(lv)\n",
    "    \n",
    "    fig.add_subplot(133)\n",
    "    sv = features['secondary_view'][0].numpy()\n",
    "    plt.plot(sv)\n",
    "    sv = ndimage.gaussian_filter(sv, 1.0)\n",
    "    plt.plot(sv)\n",
    "    \n",
    "    preds = model(features)\n",
    "    print('No filter:     ', preds[0][0].numpy().item())\n",
    "    \n",
    "    features['local_view'] = tf.expand_dims(lv, 0)\n",
    "\n",
    "    preds = model(features)\n",
    "    print('Local:         ', preds[0][0].numpy().item())\n",
    "    \n",
    "    features['global_view'] = tf.expand_dims(gv, 0)\n",
    "    \n",
    "    preds = model(features)\n",
    "    print('Global + Local:', preds[0][0].numpy().item())\n",
    "    \n",
    "    features['secondary_view'] = tf.expand_dims(sv, 0)\n",
    "    \n",
    "    preds = model(features)\n",
    "    print('All:           ', preds[0][0].numpy().item())\n",
    "    print(list(zip([\"disp_E\", \"disp_N\", \"disp_J\", \"disp_S\", \"disp_B\"], preds[0])))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
